package com.radar.example;

import cn.hutool.core.date.DateUnit;
import cn.hutool.core.date.DateUtil;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.ZipUtil;
import cn.hutool.log.Log;
import cn.hutool.log.LogFactory;
import com.radar.algorithm.DenseLucaskanade;
import com.radar.algorithm.MathCommon;
import com.radar.algorithm.SemiLagrangian;
import com.radar.common.DateConstants;
import com.radar.common.GisConstants;
import com.radar.common.ProductType;
import com.radar.exception.RadarCaclError;
import com.radar.exception.RadarDecodeError;
import com.radar.exception.RadarVisualizeError;
import com.radar.io.CinradRadar;
import com.radar.utils.BitConverter;
import com.radar.visualize.RasterMap;
import org.bytedeco.javacpp.indexer.Indexer;
import org.bytedeco.opencv.global.opencv_imgcodecs;
import org.bytedeco.opencv.global.opencv_imgproc;
import org.bytedeco.opencv.opencv_core.Mat;
import org.junit.Test;
import ucar.ma2.InvalidRangeException;


import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.util.HashSet;
import java.util.List;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

import static org.bytedeco.opencv.global.opencv_video.calcOpticalFlowPyrLK;

public class CaclExample {

    static Log log = LogFactory.get();

    @Test
    public void cappi() throws RadarDecodeError, RadarCaclError, RadarVisualizeError, IOException {
        long time = System.currentTimeMillis();
        String file = "F:/data/雷达/短临雷达数据/原始数据/Z9090/Z_RADR_I_Z9090_20190729235500_O_DOR_SB_CAP.bin";
        CinradRadar cinradRadar = new CinradRadar();
        cinradRadar.parse(file);
        double[] alt = new double[]{1000, 1500, 2000, 2500, 3000, 4000, 5000, 6000, 7000, 8000, 9000, 10000, 11000, 12000};
        byte[][][] cappi = cinradRadar.radar.cappi(ProductType.REF, alt);
        System.out.println(System.currentTimeMillis() - time);
        byte[] bytes = new byte[cappi.length * 460 * 460];
        int index = 0;
        for (int i = 0; i < cappi.length; i++) {
            for (int j = 0; j < 460; j++) {
                for (int k = 0; k < 460; k++) {
                    if (cappi[i][j * 2][k * 2] == -128 || MathCommon.DecodeRef(cappi[i][j * 2][k * 2]) <= 10) {
                        bytes[index] = (byte) 255;
                    } else {
                        bytes[index] = cappi[i][j * 2][k * 2];
                    }
                    index++;
                }
            }
        }
        File f = new File("/Users/xx/code/3d/static/3d/surface.dat");
        if (f.exists()) {
            f.delete();
        }
        FileUtil.writeBytes(bytes, f);
        ZipUtil.zip(f.getAbsolutePath(), f.getAbsolutePath() + ".zip");
        f.delete();
        RasterMap rasterMap = new RasterMap(cinradRadar.radar);
        for (int i = 0; i < cappi.length; i++) {
            BufferedImage bufferedImage = rasterMap.drawSingleRaster(MathCommon.decode(cappi[i], ProductType.REF), Color.black);
            ImageIO.write(bufferedImage, "png", new File("/Users/xx/DeskTop/cappi_" + i + ".png"));
        }
        cinradRadar.dispose();
        log.info("time : {}", System.currentTimeMillis() - time);
    }

    @Test
    public void forecast() throws RadarDecodeError, RadarCaclError, RadarVisualizeError, IOException, InvalidRangeException {
//        CinradRadar radar = new CinradRadar();
//        radar.parse("/Users/xx/DeskTop/Z9737/Z_RADR_I_Z9737_20210119181127_O_DOR_SA_CAP_FMT.bin.bz2");
//        RasterMap rasterMap = new RasterMap(radar.radar);
//
//        double[][] decode = MathCommon.decode(radar.radar.getCR().radarData, ProductType.REF);
//        for(int i = 0; i < decode.length ; i ++){
//            for(int j = 0 ; j < decode[0].length; j ++){
//                decode[i][j] = decode[i][j] != GisConstants.UNDEF && decode[i][j] < 10 ? 0 : decode[i][j];
//            }
//        }
//        decode = Cacl.nineMedianFilter(decode);
//        decode = Cacl.nineMedianFilter(decode);
//        BufferedImage bufferedImage = rasterMap.drawSingleRaster(decode, ProductType.REFGRAY);
//        ImageIO.write(bufferedImage,"png",new File("/Users/xx/DeskTop/pre1.png"));
//
//        radar = new CinradRadar();
//        radar.parse("/Users/xx/DeskTop/Z9737/Z_RADR_I_Z9737_20210119181720_O_DOR_SA_CAP_FMT.bin.bz2");
//        rasterMap = new RasterMap(radar.radar);
//
//        decode = MathCommon.decode(radar.radar.getCR().radarData, ProductType.REF);
//        for(int i = 0; i < decode.length ; i ++){
//            for(int j = 0 ; j < decode[0].length; j ++){
//                decode[i][j] = decode[i][j] != GisConstants.UNDEF && decode[i][j] < 10 ? 0 : decode[i][j];
//            }
//        }
//        decode = Cacl.nineMedianFilter(decode);
//        decode = Cacl.nineMedianFilter(decode);
//        bufferedImage = rasterMap.drawSingleRaster(decode, ProductType.REFGRAY);
//        ImageIO.write(bufferedImage,"png",new File("/Users/xx/DeskTop/pre.png"));
//
//
//        CinradRadar radar1 = new CinradRadar();
//        radar1.parse("/Users/xx/DeskTop/Z9737/Z_RADR_I_Z9737_20210119182909_O_DOR_SA_CAP_FMT.bin.bz2");
//
//        decode = MathCommon.decode(radar1.radar.getCR().radarData, ProductType.REF);
//        bufferedImage = rasterMap.drawSingleRaster(decode, ProductType.REFGRAY);
//        ImageIO.write(bufferedImage,"png",new File("/Users/xx/DeskTop/flow.png"));
//
//        for(int i = 0; i < decode.length ; i ++){
//            for(int j = 0 ; j < decode[0].length; j ++){
//                decode[i][j] = decode[i][j] != GisConstants.UNDEF && decode[i][j] < 10 ? 0 : decode[i][j];
//            }
//        }
//        decode = Cacl.nineMedianFilter(decode);
//        decode = Cacl.nineMedianFilter(decode);
//        bufferedImage = rasterMap.drawSingleRaster(decode, ProductType.REFGRAY);
//        ImageIO.write(bufferedImage,"png",new File("/Users/xx/DeskTop/next.png"));


        String f1 = "/Users/xx/DeskTop/pre.png";
        int h = ImageIO.read(new File(f1)).getHeight();
        int w = ImageIO.read(new File(f1)).getWidth();

        byte[] bytes = FileUtil.readBytes("/Users/xx/data/fy4_flow.dat");
        float[][][] V = new float[2][h][w];
        int index = 0;
        for (int i = 0; i < V.length; i++) {
            for (int j = 0; j < V[0].length; j++) {
                for (int k = 0; k < V[0][0].length; k++) {
                    V[i][j][k] = (float) BitConverter.toDouble(bytes, index);
                    index += 8;
                }
            }
        }


        String f2 = "/Users/xx/DeskTop/next.png";
        String f3 = "/Users/xx/DeskTop/flow.png";
        Mat ff1 = opencv_imgcodecs.imread(f1);
        Mat ff2 = opencv_imgcodecs.imread(f2);
//        Mat ff3 = opencv_imgcodecs.imread(f3);
//        Mat prvs = ff1.clone();
        Mat next = ff1.clone();
//        Mat rn = ff1.clone();
//        Mat flow = ff1.clone();
//        opencv_imgproc.cvtColor(ff1, prvs, opencv_imgproc.COLOR_BGR2GRAY);
        opencv_imgproc.cvtColor(ff2, next, opencv_imgproc.COLOR_BGR2GRAY);
//        opencv_imgproc.cvtColor(ff3, rn, opencv_imgproc.COLOR_BGR2GRAY);
//        // 计算光流
//        calcOpticalFlowPyrLK(prvs, next, flow, 0.5, 3, 15, 3, 5, 1.2, 0);
//        Indexer indexer = flow.createIndexer();
        Indexer indexerPrecip = next.createIndexer();

        float[][] precip = new float[h][w];
//        float[][] u = new float[h][w];
//        float[][] v = new float[h][w];
        for (int i = 0; i < h; i++) {
            for (int j = 0; j < w; j++) {
//                u[i][j] = (float) (indexer.getDouble(i, j, 0));
//                v[i][j] = (float) (indexer.getDouble(i, j, 1));
                precip[i][j] = (float) indexerPrecip.getDouble(i, j);
            }
        }
        SemiLagrangian semiLagrangian = new SemiLagrangian();
        int t = 24;
        RasterMap rasterMap = new RasterMap(null);
        float[][][] extrapolate = semiLagrangian.extrapolate(precip, V[0], V[1], t);
        for (int i = 0; i < t; i++) {
            BufferedImage bufferedImage = new BufferedImage(w, h, BufferedImage.TYPE_INT_ARGB);
            for (int j = 0; j < h; j++) {
                for (int k = 0; k < w; k++) {
                    extrapolate[i][j][k] = extrapolate[i][j][k] == 0 || extrapolate[i][j][k] == Float.NaN ? GisConstants.UNDEF : (float) (extrapolate[i][j][k] / 255.0 * 80);
                    Color color = rasterMap.getColor(extrapolate[i][j][k], ProductType.REF);
                    if (color != null) {
                        bufferedImage.setRGB(k, j, color.getRGB());
                    }
                }
            }
            ImageIO.write(bufferedImage, "png", new File("/Users/xx/data/java/flow_" + i + ".png"));
        }
    }

    @Test
    public void forecast784() throws Exception {
        String input = "/Users/xx/Downloads/img_data/";
        HashSet<String> set = new HashSet<>();
        FileUtil.loopFiles(input).stream().filter(i ->
                Pattern.matches("^(\\d{14})_VOL_(.*?).png$", i.getName()))
                .forEach(i -> set.add(i.getName().substring(0, 14)));
        List<String> files = set.stream().sorted().collect(Collectors.toList());
        for (int z = 1; z < files.size(); z++) {
            String date = files.get(z);
            String outDir = input + File.separator + date.substring(0, 8) + File.separator + "VOL" + File.separator;
            List<File> subFile = FileUtil.loopFiles(outDir).stream().filter(i ->
                    Pattern.matches("^" + date + "_VOL_LAYER_(.*?)_GRAY.png$", i.getName())).collect(Collectors.toList());
            long offset = DateUtil.between(
                    DateUtil.parse(files.get(z - 1), DateConstants.DATE_TIME_FORMAT_YYYYMMDDHHMISS),
                    DateUtil.parse(files.get(z), DateConstants.DATE_TIME_FORMAT_YYYYMMDDHHMISS),
                    DateUnit.MINUTE
            );
            for (File f : subFile) {
                String prefile = f.getName().replace(f.getName().substring(0, 14), files.get(z - 1));
                String nextImg = f.getName();
                String outDirFlow = outDir + File.separator + nextImg.substring(0, 15) + nextImg.split("_")[3] + File.separator;
                try {
                    forecastItem(outDir, prefile, nextImg, (int) offset, outDirFlow);
                }catch (Exception e){
                    e.printStackTrace();
                }
            }
        }
    }

    public void forecastItem(String outDir, String prefile, String nextImg, int offset, String outDirFlow) throws IOException {
        FileUtil.mkdir(outDirFlow);
        Mat ff1 = opencv_imgcodecs.imread(outDir + prefile);
        Mat ff2 = opencv_imgcodecs.imread(outDir + nextImg);
        Mat prvs = ff1.clone();
        Mat next = ff1.clone();
        opencv_imgproc.cvtColor(ff1, prvs, opencv_imgproc.COLOR_BGR2GRAY);
        opencv_imgproc.cvtColor(ff2, next, opencv_imgproc.COLOR_BGR2GRAY);
        int w = prvs.size().get(0);
        int h = prvs.size().get(1);
        double[][] pre = new double[h][w];
        double[][] nex = new double[h][w];
        Indexer indexer = prvs.createIndexer();
        Indexer indexer1 = next.createIndexer();
        for (int i = 0; i < h; i++) {
            for (int j = 0; j < w; j++) {
                pre[i][j] = indexer.getDouble(i, j);
                if (pre[i][j] == 255) {
                    pre[i][j] = 0;
                }
                nex[i][j] = indexer1.getDouble(i, j);
                if (nex[i][j] == 255) {
                    nex[i][j] = 0;
                }
            }
        }
        double[][][] cacl = DenseLucaskanade.cacl(new double[][][]{pre, nex});
        double[][] u = cacl[0];
        double[][] v = cacl[1];
        for (int i = 0; i < h; i++) {
            for (int j = 0; j < w; j++) {
                u[i][j] = u[i][j] * (6.0 / offset);
                v[i][j] = v[i][j] * (6.0 / offset);
            }
        }
        double[][] decode = new double[u.length][u[0].length];
        // tran 灰度图
        BufferedImage read = ImageIO.read(new File(outDir + nextImg));
        for (int i = 0; i < h; i++) {
            for (int j = 0; j < w; j++) {
                int r = new Color(read.getRGB(j, i)).getRed();
                if (r == 0) {
                    decode[i][j] = GisConstants.UNDEF;
                } else {
                    decode[i][j] = -40 + (r) / 255.0 * 60;
                }
            }
        }

        SemiLagrangian semiLagrangian = new SemiLagrangian();
        int t = 20;
        float[][][] extrapolate = semiLagrangian.extrapolate(MathCommon.doubleToFloat(decode), MathCommon.doubleToFloat(u), MathCommon.doubleToFloat(v), t);
        for (int i = 0; i < t; i++) {
            BufferedImage bufferedImage;
            bufferedImage = new BufferedImage(w, h, BufferedImage.TYPE_INT_ARGB);
            for (int j = 0; j < u.length; j++) {
                for (int k = 0; k < u[0].length; k++) {
                    if (extrapolate[i][j][k] == 0 || extrapolate[i][j][k] == Float.NaN) {
                        bufferedImage.setRGB(k, j, new Color(0, 0, 0, 0).getRGB());
                    } else if (extrapolate[i][j][k] == 20) {
                        bufferedImage.setRGB(k, j, new Color(232, 38, 38).getRGB());
                    } else {
                        Color color = GetColorByDBZ(extrapolate[i][j][k]);
                        if (color == null) {
                            bufferedImage.setRGB(k, j, new Color(0, 0, 0, 0).getRGB());
                        } else {
                            bufferedImage.setRGB(k, j, color.getRGB());
                        }
                    }
                }
            }
            ImageIO.write(bufferedImage, "png", new File(outDirFlow + File.separator + "Z_" + (i + 1) + ".png"));
        }
    }


    public static Color GetColorByDBZ(double dbz) {
        Color t_clr = null;
        if (dbz <= -40) {
            t_clr = new Color(0, 0, 0, 0);
        } else if (dbz <= -36) {
            t_clr = new Color(13, 0, 130);
        } else if (dbz <= -32) {
            t_clr = new Color(0, 0, 191);
        } else if (dbz <= -28) {
            t_clr = new Color(2, 54, 232);
        } else if (dbz <= -24) {
            t_clr = new Color(2, 97, 253);
        } else if (dbz <= -20) {
            t_clr = new Color(41, 172, 250);
        } else if (dbz <= -16) {
            t_clr = new Color(13, 220, 237);
        } else if (dbz <= -12) {
            t_clr = new Color(2, 208, 136);
        } else if (dbz <= -8) {
            t_clr = new Color(0, 191, 67);
        } else if (dbz <= -4) {
            t_clr = new Color(113, 203, 46);
        } else if (dbz <= 0) {
            t_clr = new Color(196, 224, 61);
        } else if (dbz <= 4) {
            t_clr = new Color(255, 255, 112);
        } else if (dbz <= 8) {
            t_clr = new Color(248, 205, 118);
        } else if (dbz <= 12) {
            t_clr = new Color(247, 193, 83);
        } else if (dbz <= 16) {
            t_clr = new Color(240, 163, 44);
        } else if (dbz <= 20) {
            t_clr = new Color(234, 111, 46);
        }
        return t_clr;
    }


}
